#!/usr/bin/env python3
"""
Copyright (c) [2023] [Joel Shu]
[mig-way] is licensed under Mulan PSL v2.
You can use this software according to the terms and conditions of the Mulan PSL v2. 
You may obtain a copy of Mulan PSL v2 at:
        http://license.coscl.org.cn/MulanPSL2 
THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.  
See the Mulan PSL v2 for more details.  
"""
import os
from os import path
from configparser import ConfigParser
import asyncmy
import re
import asyncio
from datetime import datetime
import logging
import sys

logging.basicConfig(
    level=logging.INFO,
    stream=sys.stdout,
    format="[%(asctime)s] [%(name)s] %(levelname)s: %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger("main")

# mig-way数据库表
MIGWAY_TABLE = """CREATE TABLE IF NOT EXISTS `{}` (
	`id` INT(11) NOT NULL AUTO_INCREMENT,
	`version` VARCHAR(50) NULL DEFAULT NULL COLLATE 'utf8_general_ci' COMMENT '版本号',
	`description` VARCHAR(200) NOT NULL COLLATE 'utf8_general_ci' COMMENT '描述',
	`script` VARCHAR(1000) NOT NULL COLLATE 'utf8_general_ci' COMMENT '文件名称',
	`installed_on` TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP COMMENT '命令运行时间',
	`execution_time` INT(11) NOT NULL DEFAULT 1 COMMENT '执行花费时间',
	`success` TINYINT(1) NOT NULL DEFAULT 1 COMMENT '状态,1-执行成功,0-执行失败',
	PRIMARY KEY (`id`) USING BTREE,
	UNIQUE INDEX `version_idx` (`version`) USING BTREE,
    INDEX `success_idx` (`success`) USING BTREE
)
COLLATE='utf8_general_ci'
COMMENT='数据库迁移版本管理表'
ENGINE=InnoDB;"""


def dict_value(data: dict, key: str, default=None):
    """获取字典的某个key,如果key不存在则返回默认值

    Args:
        key (str): 需要获取的 key
        data (dict): 字典
        default (any, optional): 如果key不存在返回的默认值. Defaults to None.

    Returns:
        any: 获取到的key对应的值
    """
    return data[key] if is_blank(data.get(key, None)) is False else default


def load_env(p=".env"):
    """从文件加载加载环境变量

    Args:
        p (str): 环境变量文件地址, 默认为: .env
    """
    env_dict = {}
    p = path.abspath(p)
    if path.exists(p):
        with open(p, 'r', encoding="utf-8") as f:
            for line in f:
                # 如果不是注释且为规范的内容行
                if not line.startswith('#') and '=' in line:
                    key, value = line.split('=', 1)
                    env_dict[key] = value.strip()
    return env_dict


def load_config(p=None):
    """加载配置文件

    Args:
        p (str, optional): 配置文件地址. Defaults to 'config.ini'.

    Returns:
        dict: 配置属性
    """
    config_dict = {}
    if p is None:
        p = "config.ini"
    if os.path.exists(p):
        conf = ConfigParser()
        conf.read(p)
        sections = conf.sections()
        if sections and len(sections) > 0:
            for section in sections:
                options = conf.options(section)
                if options and len(options) > 0:
                    for option in options:
                        key = f'{section.upper()}_{option.upper()}'
                        config_dict[key] = conf.get(section, option)
    return config_dict


def load_cfg():
    """加载配置项, 1. 首先从环境变量 '.env' 加载, 2. 其次从 `config.ini` 加载
 
    Returns:
        dict: 将解析的转换为字典
    """
    return {**load_env(), **load_config()}


def is_blank(str, ignore_whitespace=True):
    """验证字符串是否为空

    Args:
        str (str): 待验证的字符串
        ignore_whitespace (bool, optional): 是否忽略空格(包括空白字符串以及[\\r\\t\\n]之类的制表符),
            默认为true

    Returns:
        bool: bool
    """
    if str is None:
        return True
    return len(str.strip() if ignore_whitespace is True else str) == 0


async def connet_mysql(
    url=None,
    host=None,
    user=None,
    password=None,
    db=None,
    port=None,
    migway_table='',
):
    """连接 mysql

    Args:
        url (str, optional): 连接地址. Defaults to None.
        host (str, optional): 连接host. Defaults to None.
        user (str, optional): mysql登录用户. Defaults to None.
        password (str, optional): mysql登录密码. Defaults to None.
        db (str, optional): mysql数据库名称. Defaults to None.
        port (str, optional): 端口. Defaults to None.
        migway_table (str, optional): 迁移表格名称. Defaults to ''.

    Returns:
        pool: 数据库连接池
    """
    if url:  # 使用的是连接字符串
        match = re.search(r"mysql:\/\/(\w+):(\w+)@([\d\.]+):(\d+)\/(\w+)", url)
        user = match.group(1)
        password = match.group(2)
        host = match.group(3)
        port = match.group(4)
        db = match.group(5)
    if port is None:
        port = 3306
    pool = await asyncmy.create_pool(
        host=host,
        user=user,
        password=password,
        db=db,
        port=int(port),
    )

    last_versione = None
    errors = []
    async with pool.acquire() as conn:
        # 插入默认表
        async with conn.cursor() as cursor:
            migway_init = MIGWAY_TABLE.format(migway_table)
            await cursor.execute(migway_init)
        # 查询最近一条记录
        async with conn.cursor() as cursor:
            sql = f'select `version` from `{migway_table}` where `success` = 1 order by `version` desc limit 1'
            await cursor.execute(sql)
            ret = await cursor.fetchone()
            if ret is not None:
                last_versione = ret[0]
        # 查询之前执行失败的脚本
        async with conn.cursor() as cursor:
            slsql = f'select `version` from `{migway_table}` where `success` <> 1'
            await cursor.execute(slsql)
            ret = await cursor.fetchall()
            for ret_item in ret:
                errors.append(ret_item[0])
    return (pool, last_versione, errors)


def parse_sql_file(sql_path):
    """解析 sql 文件

    Args:
        sql_path (str): sql 文件路径

    Returns:
        str[]: 每一条 sql 指令
    """
    # 保存一个命令(以 ; 分隔的每一个指令)
    scripts = []
    with open(sql_path, 'r', encoding="utf-8") as f:
        sqls = f.read()
        # 以 ; 分隔的就是每一个指令
        items = sqls.split(';')
        for item in items:
            item = item.strip()
            if not is_blank(item):
                scripts.append(f'{item};')
    return scripts


async def migration(
    dirpath,
    pool,
    last_line='1',
    history_errors=[],
    schema_table='migway_schema_history',
):
    """执行数据库迁移

    Args:
        dirpath (str): 迁移文件路径
        pool (any): 数据库连接池
        last_line (str, optional): 版本号. Defaults to '1'.
        history_errors (list, optional): 之前的错误列表. Defaults to [].
        schema_table (str, optional): 迁移记录表. Defaults to 'migway_schema_history'.
    """
    dirpath = path.abspath(dirpath.strip())
    if path.exists(dirpath):
        contents = []  # 所有文件内容
        versions = []  # 所有执行的版本号列表
        for filename in os.listdir(dirpath):
            # 文件路径
            file_path = path.join(dirpath, filename)
            path_info = path.splitext(filename)
            path_name_info = path_info[0].split('__')
            # 版本
            version = path_name_info[0]
            if version.startswith('v') or version.startswith('V'):
                version = version[1:]
            if version > last_line or version in history_errors:
                # 解析内容
                scripts = parse_sql_file(file_path)
                if len(scripts) > 0:
                    contents.append((version, path_name_info[1], filename, scripts))
                    versions.append(version)
        contents.sort(key=lambda c: c[0])  # 按版本号排序
        runeds = []  # 已经执行过
        if len(versions) > 0:
            async with pool.acquire() as conn:
                # 查询当前要执行且已经执行过的版本
                async with conn.cursor() as cursor:
                    sql = f'select `version` from `{schema_table}` where `success` = 1 and `version` in (%s)'
                    in_p = ', '.join(list(map(lambda x: '%s', versions)))
                    sql = sql % in_p
                    await cursor.execute(sql, versions)
                    ret = await cursor.fetchall()
                    for ret_item in ret:
                        runeds.append(ret_item[0])

        inserts = []  # 新增
        updates = []  # 修改
        # 执行每个文件的指令
        async with pool.acquire() as conn:
            for content in contents:
                if content[0] in runeds:
                    logger.warning(f'error: {content[2]} is runed!')
                else:
                    start_time = int(datetime.now().timestamp() * 1000)
                    is_success = 0
                    async with conn.cursor() as cursor:
                        try:
                            for script in content[3]:
                                await cursor.execute(script)
                            await conn.commit()
                            is_success = 1
                            logger.info(f'{filename} run success!')
                        except Exception as e:
                            is_success = 0
                            logger.error(f'{filename} run failed!')
                            logger.error(e)
                            await conn.roll_back()
                    end_time = int(datetime.now().timestamp() * 1000)
                    time_diff = end_time - start_time
                    run_item = (
                        content[0],
                        content[1],
                        content[2],
                        datetime.now(),
                        time_diff,
                        is_success,
                    )
                    # 当前版本之前提交错误, 则是修改
                    if content[0] in history_errors:
                        updates.append(run_item)
                    else:
                        inserts.append(run_item)
        # 修改历史记录
        async with pool.acquire() as conn:
            async with conn.cursor() as cursor:
                if len(inserts) > 0:
                    in_sql = insertStatement(schema_table, [
                        'version',
                        'description',
                        'script',
                        'installed_on',
                        'execution_time',
                        'success',
                    ])
                    await cursor.executemany(in_sql, inserts)
                if len(updates):
                    up_sql = (
                        f'update `{schema_table}` set `description`=%s,`script`=%s,',
                        '`installed_on`=%s,`execution_time`=%s,`success`=%s where `version`=%s',
                    )
                    for update in updates:
                        await cursor.execute(up_sql, [
                            update[1],
                            update[2],
                            update[3],
                            update[4],
                            update[5],
                            update[0],
                        ])
                await conn.commit()


def insertStatement(table, columns):
    """组装插入的sql语句

    Args:
        table (str): 插入的数据库表名
        columns (str[]): 列的集合
        Returns:
        str: 拼接好的sql语句
    """
    states = []
    c = []
    for column in columns:
        c.append(f'`{column}`')
        states.append('%s')
    return f'INSERT INTO {table}({",".join(c)}) VALUES ({",".join(states)})'


def to_bool(bool_str):
    """将一个对象转换为bool类型

    Args:
        bool_str (any): 待转换的对象

    Returns:
        bool: 转换后的bool值
    """
    mapping = {
        None: False,
        bool: lambda x: x,
        int: lambda x: x == 1,
        str: lambda x: x.strip().lower() in ['1', 'true', '']
    }
    return mapping.get(type(bool_str), lambda x: False)(bool_str)


async def main(cfg):
    logger.info('start migration data...')
    # 迁移目录
    mig_dir = cfg.get('MIG_WAY_LOCATION')
    if is_blank(mig_dir):
        mig_dir = 'migrations'
    # 是否开启迁移
    enable = to_bool(cfg.get('MIG_WAY_ENABLED'))
    # 上一次的版本
    last_line = cfg.get('MIG_WAY_BASE_LINE')
    if enable and not is_blank(mig_dir):
        tablename = cfg.get('MIG_WAY_TABLE')
        if is_blank(tablename):
            tablename = dict_value(cfg, 'MIG_WAY_TABLE', path.basename(os.getcwd()))
        migway_table = tablename.replace('-', '_').lower()
        mig_way_table = f'{migway_table}_migway_schema_history'
        # 连接到 mysql
        init = await connet_mysql(
            url=cfg.get('MYSQL_URL'),
            host=cfg.get('MYSQL_HOST'),
            user=cfg.get('MYSQL_USER'),
            port=cfg.get('MYSQL_PORT'),
            password=cfg.get('MYSQL_PASSWORD'),
            db=cfg.get('MYSQL_DB'),
            migway_table=mig_way_table,
        )
        if not is_blank(init[1]):
            last_line = init[1]
        if is_blank(last_line):
            last_line = '1'
        if last_line.startswith('v') or last_line.startswith('V'):
            last_line = last_line[1:]
        # 执行迁移
        await migration(mig_dir, init[0], last_line, init[2], mig_way_table)
    logger.info('data migration complete!')


if __name__ == '__main__':
    cfg = load_cfg()  # 加载配置
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    loop.run_until_complete(main(cfg))
